import re

import numpy as np
import torch
from torch.utils.data import Dataset
import pandas as pd


class TitanicDataset(Dataset):
    def __init__(self, path):
        # 读取原始的数据集
        trainData = pd.read_csv(path)
        # 获得y数据集
        try:
            self.y_data = torch.Tensor(trainData["Survived"].values.astype(np.float32).reshape(-1, 1))
        except:
            pass

        items = ["PassengerId", "Survived", "Name", "Ticket", "Cabin", "Embarked"]
        # 去除无效行
        for item in items:
            try:
                trainData.drop([item], axis=1, inplace=True)
            except:
                pass
        # 设置数据集的长度
        self.len = len(trainData)
        # 遍历读到的数据进行预处理
        for i in trainData.index:
            trainData.loc[i, "Sex"] = 0 if trainData.loc[i, "Sex"] == "male" else 1
        # 获得age的中位数
        avg = trainData["Age"].mode()[0]
        # 将中位数填充nan值
        trainData["Age"] = trainData["Age"].fillna(avg)
        self.x_data = torch.Tensor(trainData.values.astype(np.float32))

    def __getitem__(self, index):
        return self.x_data[index], self.y_data[index]

    def __len__(self) -> int:
        return self.len


if __name__ == '__main__':
    titanicDataset = TitanicDataset("data/train.csv")
